from IdentityBasedEncryption import *


class BasicIdent(IdentityBasedEncryption):
    def __init__(self, lm = 16):
        """ lm is the bytes of message, 16 bytes = 128 bits """
        super(BasicIdent, self).__init__()
        self.__lm = lm
        if lm not in [16, 32]:
            raise ValueError('only support 128 bits (16 bytes) or 256 bits (32 bytes) message')
        else:
            self.__hash = md5 if lm == 16 else sha256

    def __h2(self, gt):
        """ GT -> {0, 1} ^ lm """
        bgt = str(gt).encode('utf-8')
        return bytes.fromhex(self.__hash(bgt).hexdigest())

    def system_build(self):
        s = self.group.random(ZR)
        return {'msk': s, 'mpk': s * self.g}

    def key_extract(self, user_id, system_params):
        Q_U = self.from_hash(user_id, G1, hash_algorithm = lambda m: self.__hash(m).hexdigest())
        return {'usk': system_params['msk'] * Q_U, 'upk': Q_U}

    def encryption(self, m, system_params, user_params):
        if len(m) != self.__lm:
            raise ValueError(f'the bytes of m not satisfied the required length {self.__lm}')
        r = self.group.random(ZR)
        return r * self.g, bxor(m, self.__h2(e(system_params['mpk'], user_params['upk']) ** r))

    def decryption(self, ciphertext, user_params):
        V, W = ciphertext
        return bxor(W, self.__h2(e(V, user_params['usk'])))

if __name__ == '__main__':
    bi = BasicIdent(16)
    bi.test(b'0123456789zbcdef', b'xxx')
